from langchain_core.prompts import PromptTemplate
from langchain_milvus import Milvus
from langchain_ollama import OllamaEmbeddings, OllamaLLM
# 初始化 LLM 和 Embeddings
llm = OllamaLLM(model="qwen2.5:14b", base_url="http://192.168.7.3:11434")
embedder = OllamaEmbeddings(model="nomic-embed-text:latest", base_url="http://192.168.7.3:11434")

# 连接 Milvus 向量数据库
vectorstore = Milvus(
    embedding_function=embedder,
    collection_name="knowledge",
    connection_args={"uri": "http://192.168.6.20:19530/ipp_air_general"},
    auto_id=True
)

retriever = vectorstore.as_retriever()

# **Step 1: LLM 进行查询扩展**
query_expansion_prompt = PromptTemplate(
    input_variables=["query"],
    template="请基于用户的问题 `{query}` 生成 3 个不同的查询方式，以便更好地进行信息检索。\n每个查询请用换行分隔。"
)

expanded_queries = llm.predict(query_expansion_prompt.format(query="基本污染源包括哪些?"))

# 将 LLM 生成的查询转换为列表
expanded_queries = expanded_queries.strip().split("\n")

# **Step 2: 进行向量检索**
retrieved_docs = []
for q in expanded_queries:
    docs = retriever.get_relevant_documents(q)
    retrieved_docs.extend(docs)  # 合并所有查询的检索结果

# **Step 3: 使用 LLM 进行结果总结**
context = "\n".join([doc.page_content for doc in retrieved_docs])

summary_prompt = PromptTemplate(
    input_variables=["context"],
    template="基于以下信息，请总结回答‘基本污染源包括哪些?’\n\n{context}\n\n总结："
)

summary = llm.predict(summary_prompt.format(context=context))

# **Step 4: 输出最终答案**
print("【查询扩展】")
for i, q in enumerate(expanded_queries):
    print(f"{i+1}. {q}")

print("\n【最终回答】")
print(summary)